package core

import (
	"StressTestOpenAI/core/model"
	"fmt"
	"github.com/xuri/excelize/v2"
	"os"
	"sync"
	"time"
)

// 报告详情
type Report struct {
	Current           int
	Errs              int64
	FirstRuestTime    int64
	FirstTokenTime    int64
	TokenPS           string
	TokenAvgPS        string
	TotalTime         string
	TotalRequests     int64   // 总请求数
	TotalInputTokens  int64   // 总输入token
	TotalOutputTokens int64   // 总输出token
	TotalTokenPS      float64 // 总输出token吞吐（token/s）
	AvgTokens         int64
	EndTokenTime      int64
	TTFT_Total        int64   // 总的首token响应时间
	TTFT_Avg          string  // 平均首token响应时间
	TPOT_Total        float64 // (token结束时间-首token时间) / (tokens-1)
	TPOT_Avg          string
	ITL_Total         float64 // (token结束时间-请求开始时间) / (tokens-1)
	ITL_Avg           string
	E2EL_Total        int64
	E2EL_Avg          string

	ContentLenTotal int64 // 输出字符串的总长度
	ContentLenAvg   string
	f               *excelize.File
	Metrics         []*model.Metric
}

func NewReport(current int) *Report {
	return &Report{
		Current: current,
		f:       excelize.NewFile()}
}

func (report *Report) nitFile() {
	report.f = excelize.NewFile()
}

// 创建excel的表
func (report *Report) NewFileSheet() (indexTotal, indexDetail int, err error) {
	indexTotal, err = report.f.NewSheet("总况")
	if err != nil {
		fmt.Println("新建总况表错误：", err.Error())
		return
	}
	indexDetail, err = report.f.NewSheet("明细")
	if err != nil {
		fmt.Println("新建明细表错误：", err.Error())
	}
	return
}

func (report *Report) writeData() {
	defer func() {
		if err := report.f.Close(); err != nil {
			fmt.Println("关闭excel失败：  ", err.Error())
		}
	}()
	indexTotal, indexDetail, err := report.NewFileSheet()
	if err != nil {
		return
	}
	report.writeTotalData(indexTotal)
	report.writeDetailData(indexDetail)

	path, err := os.Getwd()
	if err != nil {
		fmt.Println("获取当前路径失败：  ", err.Error())
		path = "."
	}
	now := time.Now()
	file_name := fmt.Sprintf("%s\\report\\并发_%v_%v_%v_%v.xlsx", path, report.Current, now.Format("2006-01-02"), now.Format("15"), now.Unix())
	if _, err = os.Stat(file_name); os.IsExist(err) {
		err = os.Remove(file_name)
		if err != nil {
			fmt.Println("删除文件失败：  ", err.Error())
			return
		}
	}
	err = report.f.SaveAs(file_name)
	if err != nil {
		fmt.Println("保存文件失败：  ", err.Error())
	}
}

var total_header_list = []string{
	"总请求数", "并发数", "输出总token", "输入总token", "总耗时", "输出字符串的总长度",
	"平均每个token的字符串长度", "总输出吞吐(tokens/s)", "每个请求平均token(tokens/个)",
	"TTFT_Avg(ms)", "TPOT_Avg(ms)", "ITL_Avg(ms)", "E2EL_Avg(ms)", "错误数",
}

func (report *Report) writeTotalData(index int) {
	// 设置表头
	/*
			"总请求数   " +
			"并发数      " +
			"输出总token    " +
			"输入总token     " +
			"总耗时      " +
			"输出字符串的总长度       " +
			"平均每个token的字符串长度       " +
			"总输出吞吐(tokens/s)      " +
			"每个请求平均token(tokens/个)       " +
			"TTFT_Avg(ms)       " +
			"TPOT_Avg(ms)       " +
			"ITL_Avg(ms)        " +
			"E2EL_Avg(ms)     " +
			"错误数"

	**/

	sheet := report.f.GetSheetName(index)
	for i, h := range total_header_list {
		cell := fmt.Sprintf("%s1", string(byte(i+65)))
		report.writeValue(sheet, cell, h)
	}

	report.writeValue(sheet, "A2", report.TotalRequests)
	report.writeValue(sheet, "B2", report.Current)
	report.writeValue(sheet, "C2", report.TotalOutputTokens)
	report.writeValue(sheet, "D2", report.TotalInputTokens)
	report.writeValue(sheet, "E2", report.TotalTime)
	report.writeValue(sheet, "F2", report.ContentLenTotal)
	report.writeValue(sheet, "G2", report.ContentLenAvg)
	report.writeValue(sheet, "H2", report.TokenPS)
	report.writeValue(sheet, "I2", report.TokenAvgPS)
	report.writeValue(sheet, "J2", report.TTFT_Avg)
	report.writeValue(sheet, "K2", report.TPOT_Avg)
	report.writeValue(sheet, "L2", report.ITL_Avg)
	report.writeValue(sheet, "M2", report.E2EL_Avg)
	report.writeValue(sheet, "N2", report.Errs)

	report.f.SetActiveSheet(index)
}

func (report *Report) writeValue(sheet, cell string, value interface{}) {
	if err := report.f.SetCellValue(sheet, cell, value); err != nil {
		fmt.Println(fmt.Sprintf("表格： %v, 行列：%v, 写入数据：%v,  失败： ", sheet, cell, value, err.Error()))
	}
}

var detail_header_list = []string{"输入token数", "输出token数", "开始请求时间", "首token时间", "输出结束时间"}

func (report *Report) writeDetailData(index int) {
	if len(report.Metrics) < 1 {
		return
	}
	sheet := report.f.GetSheetName(index)
	for i, h := range detail_header_list {
		cell := fmt.Sprintf("%s1", string(byte(i+65)))
		report.writeValue(sheet, cell, h)
	}
	for j, metric := range report.Metrics {
		x := j + 2
		cell := fmt.Sprintf("A%v", x)
		report.writeValue(sheet, cell, fmt.Sprintf("%v", metric.InputTokens))
		cell = fmt.Sprintf("B%v", x)
		report.writeValue(sheet, cell, fmt.Sprintf("%v", metric.OutputTokens))
		cell = fmt.Sprintf("C%v", x)
		report.writeValue(sheet, cell, fmt.Sprintf("%v", metric.RequestTime))
		cell = fmt.Sprintf("D%v", x)
		report.writeValue(sheet, cell, fmt.Sprintf("%v", metric.FirstTokenTime))
		cell = fmt.Sprintf("E%v", x)
		report.writeValue(sheet, cell, fmt.Sprintf("%v", metric.EndTokenTime))
	}
	report.f.SetActiveSheet(index)
}
func RunReport(current int, wg *sync.WaitGroup, channel chan *model.Metric) {
	defer wg.Done()
	report := NewReport(current)
	stop := make(chan bool, 1)
	wg.Add(1)
	go printReport(report, stop, wg)
	for metric := range channel {
		report.TotalRequests++
		report.Metrics = append(report.Metrics, metric)
		if metric.Err != "" {
			report.Errs++
			continue
		}

		if report.FirstRuestTime == 0 || report.FirstRuestTime > metric.RequestTime {
			report.FirstRuestTime = metric.RequestTime
		}
		if report.FirstTokenTime == 0 || report.FirstTokenTime > metric.FirstTokenTime {
			report.FirstTokenTime = metric.FirstTokenTime
		}

		if report.EndTokenTime == 0 || report.EndTokenTime < metric.EndTokenTime {
			report.EndTokenTime = metric.EndTokenTime
		}

		report.TTFT_Total += metric.TTFT
		report.TPOT_Total += metric.TPOT
		report.ITL_Total += metric.ITL
		report.E2EL_Total += metric.E2EL
		report.TotalInputTokens += metric.InputTokens
		report.TotalOutputTokens += metric.OutputTokens
		report.ContentLenTotal += int64(metric.ContentLen)
		report.ContentLenAvg = fmt.Sprintf("%0.2f", float64(report.ContentLenTotal)/float64(report.TotalOutputTokens))

		success_total := report.TotalRequests - report.Errs
		report.TTFT_Avg = fmt.Sprintf("%0.2f", float64(report.TTFT_Total)/float64(success_total))
		report.TPOT_Avg = fmt.Sprintf("%0.2f", float64(report.TPOT_Total)/float64(success_total))
		report.ITL_Avg = fmt.Sprintf("%0.2f", float64(report.ITL_Total)/float64(success_total))
		report.E2EL_Avg = fmt.Sprintf("%0.2f", float64(report.E2EL_Total)/float64(success_total))
		totalTime := float64(report.EndTokenTime-report.FirstRuestTime) / 1000
		report.TotalTime = fmt.Sprintf("%0.2f", totalTime)
		report.TokenPS = fmt.Sprintf("%0.2f", float64(report.TotalOutputTokens)/totalTime)
		report.TokenAvgPS = fmt.Sprintf("%0.2f", float64(report.TotalOutputTokens)/float64(success_total))
	}
	stop <- true
}

func printReport(report *Report, stop chan bool, wg *sync.WaitGroup) {
	ticker := time.NewTicker(10 * time.Second)
	defer wg.Done()
	defer ticker.Stop()
	fmt.Println(
		"总请求数   " +
			"并发数      " +
			"输出总token    " +
			"输入总token     " +
			"总耗时      " +
			"输出字符串的总长度       " +
			"平均每个token的字符串长度       " +
			"总输出吞吐(tokens/s)      " +
			"每个请求平均token(tokens/个)       " +
			"TTFT_Avg(ms)       " +
			"TPOT_Avg(ms)       " +
			"ITL_Avg(ms)        " +
			"E2EL_Avg(ms)     " +
			"错误数")
	for {
		select {
		case <-ticker.C:
			if report.TotalRequests == 0 {
				continue
			}
			fmt.Println(fmt.Sprintf("   %v      %v             %v           %v          %v              %v                     %v                     %v                   %v                        %v               %v            %v              %v         %v",
				report.TotalRequests,
				report.Current,
				report.TotalOutputTokens,
				report.TotalInputTokens,
				report.TotalTime,
				report.ContentLenTotal,
				report.ContentLenAvg,
				report.TokenPS,
				report.TokenAvgPS,
				report.TTFT_Avg,
				report.TPOT_Avg,
				report.ITL_Avg,
				report.E2EL_Avg,
				report.Errs))

		case <-stop:
			if report.TotalRequests == 0 {
				continue
			}
			fmt.Println(fmt.Sprintf("   %v      %v             %v           %v          %v              %v                     %v                     %v                   %v                        %v               %v            %v              %v         %v",
				report.TotalRequests,
				report.Current,
				report.TotalOutputTokens,
				report.TotalInputTokens,
				report.TotalTime,
				report.ContentLenTotal,
				report.ContentLenAvg,
				report.TokenPS,
				report.TokenAvgPS,
				report.TTFT_Avg,
				report.TPOT_Avg,
				report.ITL_Avg,
				report.E2EL_Avg,
				report.Errs))
			report.writeData()
			return
		}
		time.Sleep(1 * time.Second)
	}
}
